import os
import tensorflow as tf
from tensorflow import keras
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt

def plot_image(data, prediction_array):
  img = data
  plt.grid(False)
  plt.xticks([])
  plt.yticks([])

  plt.imshow(img, cmap=plt.cm.binary)

  predicted_label = np.argmax(prediction_array)

  plt.xlabel("{} {:2.0f}%".format(predicted_label,
                                100*np.max(prediction_array),),
                                color='blue')

def plot_value_array(prediction_array):
  plt.grid(False)
  plt.xticks(range(10))
  plt.yticks([])
  thisplot = plt.bar(range(10), prediction_array, color="#777777")
  plt.ylim([0, 1])
  predicted_label = np.argmax(prediction_array)

  thisplot[predicted_label].set_color('blue')

print("加载模型")
model = tf.keras.models.load_model('./mnist_checkpoint/')
probability_model = tf.keras.Sequential([model,tf.keras.layers.Softmax()])
print("加载完成")
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

while True:
    png_path = "./db/num.png"
    png = Image.open(png_path)
    png = png.convert("L")
    dt = np.zeros((28, 28), dtype=int)
    # print(dt)
    for y in range(png.size[1]):
        for x in range(png.size[0]):
            pixel = png.getpixel((x, y))
            dt[y][x] = 255 - pixel

    # dt = dt / 255.0
    # print(dt)
    # print(dt, train_images[0])
    # input()
    predictions = probability_model.predict(np.array([dt]))
    plt.figure(figsize=(6,3))
    plt.subplot(1,2,1)
    plot_image(dt, predictions[0])
    plt.subplot(1,2,2)
    plot_value_array(predictions[0])
    plt.show()
